import os
import glob

# 设置路径B
path_B = r'C:\Users\Caibin\Desktop\mixup_20250123_seg\mixup\train\images'  # 替换为你的图片文件夹路径
output_file = r'C:\Users\Caibin\Desktop\mixup_20250123_seg\mixup\train\train.txt'  # 输出文件名


def get_image_files(directory):
    # 支持的图片格式
    image_extensions = ["*.jpg", "*.jpeg", "*.png", "*.bmp", "*.gif", "*.tiff", "*.webp"]
    image_files = []

    # 获取所有图片文件的完整路径（包括子文件夹）
    for ext in image_extensions:
        image_files.extend(glob.glob(os.path.join(directory, "**", ext), recursive=True))

    return image_files


def main():
    # 获取图片文件列表
    image_files = get_image_files(path_B)

    # 将文件名写入 output.txt
    with open(output_file, 'w') as f:
        for image_file in image_files:
            # 获取第一个父文件夹和文件名
            parent_folder = os.path.basename(os.path.dirname(image_file))
            filename = os.path.basename(image_file)
            f.write(f"{parent_folder}/{filename}\n")

    print(f"已完成，图片文件名已保存至: {output_file}")


if __name__ == "__main__":
    main()
